
import threading
import Jetson.GPIO as GPIO
class PID:
    def __init__(self, P, I, D):
        self.Kp = P
        self.Ki = I
        self.Kd = D
        self.previous_error = 0
        self.integral = 0
 
    def compute(self, set_point, measured_value):
        error = measured_value - set_point
        self.integral += error
        derivative = error - self.previous_error
        output = self.Kp * error + self.Ki * self.integral + self.Kd * derivative
        self.previous_error = error
        return output
 
    def update_pid(self, P, I, D):
        self.Kp, self.Ki, self.Kd = P, I, D
 
    def get_p(self): return self.Kp
    def get_i(self): return self.Ki
    def get_d(self): return self.Kd
    
    def set_p(self, P): self.Kp = P
    def set_i(self, I): self.Ki = I
    def set_d(self, D): self.Kd = D
 
class ServoController:
    def __init__(self):
        GPIO.setmode(GPIO.BOARD)
        GPIO.setwarnings(False)
        
        # 初始化PWM通道
        self.pan_pwm = GPIO.PWM(PWM_1, 50)
        self.tilt_pwm = GPIO.PWM(PWM_2, 50)
        
        # 启动PWM（初始占空比0）
        self.pan_pwm.start(7.5)
        self.tilt_pwm.start(7.5)
 
    def set_angle(self, servo, angle):
        """设置舵机到指定角度（自动完成角度到占空比的转换）"""
        angle = max(0, min(180, angle))  # 限制角度范围
        duty_cycle = 2.5 + (angle / 180) * 10  # 线性映射到占空比
        duty_cycle = duty_cycle / 100  # 转换为小数形式
        servo.ChangeDutyCycle(duty_cycle)
        # print(f"舵机设置角度: {angle}度, 占空比: {duty_cycle:.3f}")
 
class YuntaiTrack:
    def __init__(self):
        self.servo_controller = ServoController()
        self.pan_angle = 90
        self.tilt_angle = 90
        self.pan_pid = PID(0.022, 0.0, 0.013)
        self.tilt_pid = PID(0.022, 0.0, 0.013)
        
        # 初始化舵机到中位
        self.servo_controller.set_angle(self.servo_controller.pan_pwm, self.pan_angle)
        self.servo_controller.set_angle(self.servo_controller.tilt_pwm, self.tilt_angle)
 
    def track(self, target_x, target_y, img_width, img_height):
        # 计算目标与图像中心的偏差
        error_x = target_x - (img_width / 2)
        error_y = target_y - (img_height / 2)
 
        # PID算法计算舵机调整角度
        pan_adjust = self.pan_pid.compute(0, error_x)
        tilt_adjust = self.tilt_pid.compute(0, error_y)
 
        # 更新舵机角度
        new_pan_angle = self.pan_angle + pan_adjust
        new_tilt_angle = self.tilt_angle - tilt_adjust
        
        # 限制角度范围
        new_pan_angle = max(min(new_pan_angle, 180), 0)
        new_tilt_angle = max(min(new_tilt_angle, 180), 60)
 
        # 更新角度并设置舵机
        self.pan_angle = new_pan_angle
        self.tilt_angle = new_tilt_angle
        self.servo_controller.set_angle(self.servo_controller.pan_pwm, self.pan_angle)
        self.servo_controller.set_angle(self.servo_controller.tilt_pwm, self.tilt_angle)
 
class YuntaiThread(threading.Thread):
    def __init__(self, yuntai_track, img_width, img_height):
        super().__init__()
        self.yuntai_track = yuntai_track
        self.img_width = img_width
        self.img_height = img_height
        self.target_x = None
        self.target_y = None
        self.target_name = None
        self.lock = threading.Lock()
        self.condition = threading.Condition()
        self.stop_event = threading.Event()
 
    def set_target_name(self, target_name):
        with self.lock:
            self.target_name = target_name
 
    def handle_infer_result(self, infer_result):
        if self.target_name is None:
            return
        
        boxs = infer_result["boxs"]
        classid = infer_result["classid"]
        if not boxs:
            return
        
        for i in range(len(boxs)):
            if categories[int(classid[i])] == self.target_name:
                x1, y1, x2, y2 = boxs[i]
                self.target_x = (x1 + x2) / 2
                self.target_y = (y1 + y2) / 2
                with self.condition:
                    self.condition.notify()
                break
        else:
            self.target_x = None
            self.target_y = None
 
    def run(self):
        while not self.stop_event.is_set():
            with self.condition:
                self.condition.wait()
            if self.target_x is not None and self.target_y is not None:
                self.yuntai_track.track(self.target_x, self.target_y, self.img_width, self.img_height)
 
    def stop(self):
        self.stop_event.set()
        with self.condition:
            self.condition.notify()